using System;
using System.Collections.Generic;
using System.Linq;

using Atomic.Libraries;
using Atomic.Libraries.Mathematics;
using Atomic.Libraries.Mathematics.Functions;
using Atomic.Libraries.Mathematics.LinearAlgebra;
using Atomic.Libraries.Mathematics.Optimization.Ipopt;
using Atomic.Libraries.Mathematics.Statistics;
using Atomic.Structures;

namespace Atomic.Clusters.Fitting
{
	/// <summary>
	/// Direct inversion of cluster functions. Matrix pseudo-inversion is used, so the number of structures and clusters don't need to agree.
	/// </summary>
	[Serializable]
	public class DirectInversionClusterExpansion : ClusterExpansion
	{
		private DirectInversionClusterExpansion(Structure lattice, IEnumerable<ClusterExpansionTerm> terms, IEnumerable<StructureMapping> structures, IClusterFunctionEvaluator evaluator)
			: base(lattice, terms)
		{
			Structures = new ReadOnlyList<StructureMapping>(structures);

			// Determine fitting error.
			RootMeanSquareStatistics fe = new RootMeanSquareStatistics();
			RootMeanSquareStatistics rfe = new RootMeanSquareStatistics();
			foreach (StructureMapping structure in structures)
			{
				double e0 = structure.Value;
				double e = ExpandEnergy(structure.Structure, evaluator);
				fe.Add(e - e0);
				rfe.Add(e / e0 - 1.0);
			}

			FittingError = fe.RootMeanSquare;
			RelativeFittingError = rfe.RootMeanSquare;
		}

		public static DirectInversionClusterExpansion Fit(Structure lattice, IEnumerable<StructureMapping> structures, IEnumerable<ClusterClass> clusters)
		{
			return Fit(lattice, structures, clusters, new ClusterFunctionEvaluator(lattice));
		}

		public static DirectInversionClusterExpansion Fit(Structure lattice, IEnumerable<StructureMapping> structures, IEnumerable<ClusterClass> clusters, IClusterFunctionEvaluator evaluator)
		{
			if (lattice == null || structures == null || clusters == null || evaluator == null)
			{
				throw new ArgumentNullException();
			}

			return Fit(lattice, structures.ToArray(), clusters.ToArray(), evaluator);
		}

		private static DirectInversionClusterExpansion Fit(Structure lattice, StructureMapping[] structures, ClusterClass[] clusters, IClusterFunctionEvaluator evaluator)
		{
			int n = structures.Length;
			int m = clusters.Length;

			if (n == 0 || m == 0)
			{
				throw new ArgumentException("At least one structure and one cluster are required.");
			}

			for (int i = 0; i < m; i++)
			{
				if (clusters[i].Lattice != lattice)
				{
					throw new ArgumentException("Clusters defined on different lattices.");
				}
			}

			// Evaluate correlation function for all structures at once.
			IClusterFunctionValue[] values = evaluator.Evaluate(structures.Select(s => s.Structure), clusters).ToArray();

			double[,] phi = new double[n, m];
			double[] b = new double[n];
			for (int i = 0; i < n; i++)
			{
				int j = 0;
				foreach (double c in values[i].Correlations)
				{
					phi[i, j] = clusters[j].Multiplicity * values[i].Scaling * c;
					j++;
				}

				b[i] = structures[i].Value;
			}
			Vector coefficients = SingularValueDecomposition.PseudoInverse(new Matrix(phi)) * new Vector(b);

			List<ClusterExpansionTerm> terms = new List<ClusterExpansionTerm>();
			for (int i = 0; i < m; i++)
			{
				// Factor out the multiplicity as the usual convention.
				terms.Add(new ClusterExpansionTerm(clusters[i], coefficients[i]));
			}

			return new DirectInversionClusterExpansion(lattice, terms, structures, evaluator);
		}

		public static DirectInversionClusterExpansion FitRelative(Structure lattice, IEnumerable<StructureMapping> structures, IEnumerable<ClusterClass> clusters)
		{
			return FitRelative(lattice, structures, clusters, new ClusterFunctionEvaluator(lattice));
		}

		public static DirectInversionClusterExpansion FitRelative(Structure lattice, IEnumerable<StructureMapping> structures, IEnumerable<ClusterClass> clusters, IClusterFunctionEvaluator evaluator)
		{
			if (lattice == null || structures == null || clusters == null || evaluator == null)
			{
				throw new ArgumentNullException();
			}

			return FitRelative(lattice, structures.ToArray(), clusters.ToArray(), evaluator);
		}

		private static DirectInversionClusterExpansion FitRelative(Structure lattice, StructureMapping[] structures, ClusterClass[] clusters, IClusterFunctionEvaluator evaluator)
		{
			// Use an absolute error fitting as the starting point.
			DirectInversionClusterExpansion ce = Fit(lattice, structures, clusters, evaluator);

			int n = structures.Length;
			int m = clusters.Length;

			// m coefficients to be determined.
			Variable[] coefficients = new Variable[m];
			List<VariableAssignment> assignments = new List<VariableAssignment>();
			for (int i = 0; i < m; i++)
			{
				coefficients[i] = new Variable();
				assignments.Add(new VariableAssignment(coefficients[i], ce.Terms[i].Coefficient));
			}

			// Evaluate correlation function for all structures at once.
			IClusterFunctionValue[] values = evaluator.Evaluate(structures.Select(s => s.Structure), clusters).ToArray();

			// Objective function.
			Function f = 0.0;

			for (int i = 0; i < n; i++)
			{
				int j = 0;
				Function e = 0.0;
				foreach (double c in values[i].Correlations)
				{
					e += clusters[j].Multiplicity * values[i].Scaling * c * coefficients[j];
					j++;
				}

				double e0 = structures[i].Value;
				f += Function.Sqr(e / e0 - 1.0);
			}

			f = Function.Sqrt(f / n);

			IpoptOptimizer o = new IpoptOptimizer();
			o.ObjectiveFunction = f;
			o.Variables.Add(coefficients);
			IpoptOptimizerResult or = o.RunIpopt(new Point(assignments));

			if (!or.Status)
			{
				throw new Exception("Optimization didn't converge.");
			}

			List<ClusterExpansionTerm> terms = new List<ClusterExpansionTerm>();
			for (int i = 0; i < m; i++)
			{
				// Factor out the multiplicity as the usual convention.
				terms.Add(new ClusterExpansionTerm(clusters[i], or.OptimalPoint[coefficients[i]]));
			}

			return new DirectInversionClusterExpansion(lattice, terms, structures, evaluator);
		}

		public ReadOnlyList<StructureMapping> Structures
		{
			get;
			private set;
		}

		/// <summary>
		/// The root-mean-square error (eV).
		/// </summary>
		public double FittingError
		{
			get;
			private set;
		}

		/// <summary>
		/// The relative fitting error.
		/// </summary>
		public double RelativeFittingError
		{
			get;
			private set;
		}
	}
}
